package org.example.doctorapi.config;


import dev.langchain4j.code.graalvm.GraalVmJavaScriptExecutionEngine;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.image.ImageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.model.zhipu.ZhipuAiEmbeddingModel;
import dev.langchain4j.model.zhipu.ZhipuAiImageModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import dev.langchain4j.web.search.WebSearchTool;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import org.example.doctorapi.server.*;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;
import java.util.List;

@Configuration
public class LLMConfig {

    @Bean("ChatLanguageModel1")
    public ChatLanguageModel chatLanguageModel() {
        return OpenAiChatModel.builder()
                .apiKey("sk-3bfcbc6929c04a3e9039ced3bc77f21b")
                .modelName("qwen-max")
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .maxRetries(5)
                .logRequests(true)
                .logResponses(true)
                .timeout(Duration.ofSeconds(20))
                .listeners(List.of(new TestChatModelListener()))
                .build();
    }

    @Bean("StreamingChatLanguageModel")
    public StreamingChatLanguageModel streamingChatLanguageModel(){
        return OpenAiStreamingChatModel.builder()
                .apiKey("sk-3bfcbc6929c04a3e9039ced3bc77f21b")
                .modelName("qwen-max")
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .logRequests(true)
                .logResponses(true)
                .timeout(Duration.ofSeconds(20))
                .listeners(List.of(new TestChatModelListener()))
                .build();
    }

    @Bean("StreamingChatAssistant")
    public ChatAssistantFlux chatAssistantStream(){
        return AiServices.builder(ChatAssistantFlux.class)
                .streamingChatLanguageModel(streamingChatLanguageModel())
                .build();
    }

    @Bean("ChatAssistant")
    public ChatAssistant chatAssistant() {
        return AiServices.builder(ChatAssistant.class)
                .chatLanguageModel(chatLanguageModel())
                .build();
    }

    @Bean("ChatLanguageModelImage")
    public ChatLanguageModel chatLanguageModelImage() {
        return OpenAiChatModel.builder()
                .apiKey("sk-3bfcbc6929c04a3e9039ced3bc77f21b")
                .modelName("qwen-vl-max")
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .logRequests(true)
                .logResponses(true)
                .timeout(Duration.ofSeconds(20))
                .listeners(List.of(new TestChatModelListener()))
                .build();
    }

    /*智普Ai*/
    @Bean("ZhiPuImageModel")
    public ImageModel imageModel(){
        ImageModel build = ZhipuAiImageModel.builder()
                .apiKey("2676fc8e80844f7a85b0ff5a10fc5083.1D4a1gQ38W46uCon")
                .logRequests(true)
                .logResponses(true)
                .connectTimeout(Duration.ofSeconds(10000))
                .callTimeout(Duration.ofSeconds(10000))
                .readTimeout(Duration.ofSeconds(10000))
                .writeTimeout(Duration.ofSeconds(10000))
                .build();
        return build;
    }

    @Bean("ZhipuAiEmbeddingModel")
    public ZhipuAiEmbeddingModel embeddingModel(){
         return ZhipuAiEmbeddingModel.builder()
                .apiKey("2676fc8e80844f7a85b0ff5a10fc5083.1D4a1gQ38W46uCon")
                .logRequests(true)
                .logResponses(true)
                .connectTimeout(Duration.ofSeconds(10000))
                .callTimeout(Duration.ofSeconds(10000))
                .readTimeout(Duration.ofSeconds(10000))
                .writeTimeout(Duration.ofSeconds(10000))
                .build();
    }

    @Bean("chatLanguageModelMemory")
    public ChatLanguageModel chatLanguageModelMemory(){
        return OpenAiChatModel.builder()
                .apiKey("sk-3bfcbc6929c04a3e9039ced3bc77f21b")
                .modelName("qwen-long") // 设置使用的模型名称
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .logRequests(true)
                .build();
    }

    @Bean("chatLanguageModelMemoryBean")
    public ChatAssistant chatLanguageModelMemoryBean(){
        return AiServices.builder(ChatAssistant.class)
                .chatLanguageModel(chatLanguageModelMemory())
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                .build();
    }


    @Bean
    public ChatMemoryStore chatMemoryStore() {
        return new PersistentChatMemoryStore();
    }
    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
        return memoryId -> MessageWindowChatMemory.builder()
                // 设置消息窗口 ID
                .id(memoryId)
                // 设置消息最大条数，默认为 10
                .maxMessages(10)
                .chatMemoryStore(chatMemoryStore())
                .build();
    }

    @Bean("StreamingChatAssistantMemoryDemo")
    public ChatAssistantFlux chatAssistantStreamMemoryDmo(){
        return AiServices.builder(ChatAssistantFlux.class)
                .streamingChatLanguageModel(streamingChatLanguageModel())
                .chatMemoryProvider(chatMemoryProvider())
                .build();
    }


    @Bean("AiChatAssistant")
    public ChatAssistant AichatAssistant() {
        return AiServices.builder(ChatAssistant.class)
                .chatLanguageModel(chatLanguageModel())
                /*.tools(new NumberHandler())*/
                .tools(new GraalVmJavaScriptExecutionEngine())
                .build();
    }

    /*searchWeb*/
    @Bean("searchChatAssistant")
    public SearchChatAssistant searchChatAssistant(){
        SearchApiWebSearchEngine searchEngine = SearchApiWebSearchEngine.builder()
                .apiKey("fKUCGU7YAuKDUkHXgRPruver")// 测试使用
                .engine("google")
                .build();
        WebSearchTool webSearchTool = WebSearchTool.from(searchEngine);
        return AiServices.builder(SearchChatAssistant.class)
                .chatLanguageModel(chatLanguageModel())
                .tools(webSearchTool)
                .build();
    }
    // RAG
    @Bean("MyInMemoryEmbeddingStore")
    public InMemoryEmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    @Bean("chatAssistantStreamRAG")
    public ChatAssistantFlux chatAssistantStreamRAG(){
        return AiServices.builder(ChatAssistantFlux.class)
                .streamingChatLanguageModel(streamingChatLanguageModel())
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore()))
                .build();
    }




}
